package problems.daily.y2021m7;

import java.util.HashMap;
import java.util.Map;

/**
 * @author habitplus
 * @since 2021-07-08 09:26
 */
public class T20210708N1 {
    // 930. 和相同的二元子数组
    // 给你一个二元数组 nums ，和一个整数 goal ，请你统计并返回有多少个和为 goal 的 非空 子数组。
    // 子数组 是数组的一段连续部分。
    //
    // 提示：
    //      1 <= nums.length <= 3 * 104
    //      nums[i] 不是 0 就是 1
    //      0 <= goal <= nums.length
    //
    // 链接：https://leetcode-cn.com/problems/binary-subarrays-with-sum/
    public int numSubarraysWithSum01(int[] nums, int goal) {
        if (nums == null || nums.length == 0 || goal > nums.length) {
            return 0;
        }

        int n = nums.length;
        int[] onePoses = new int[n];
        int oneCnt = 0;
        for (int i = 0; i < n; i++) {
            if (nums[i] == 1) {
                onePoses[oneCnt] = i;
                oneCnt++;
            }
        }

        int cnt = 0;
        int prev = -1;
        int leftZeroCnt, rightZeroCnt;
        for (int i = 0; i + goal - 1 < oneCnt; i++) {
            // 统计左边零的个数
            leftZeroCnt = (i < oneCnt ? onePoses[i] : n) - prev - 1;

//            cnt += (leftZeroCnt + 1) * leftZeroCnt / 2;

            // 统计右边零的个数：必须要 goal > 0 才能统计右边的零
            rightZeroCnt = 0;
            if (goal > 0) {
                // 不是最后一个 1 的位置，取 onePoses[i + goal]
                // 已经是最后一个 1 的位置，取 n
                rightZeroCnt = (i + goal < oneCnt ? onePoses[i + goal] : n) - onePoses[i + goal - 1] - 1;
//                cnt += (rightZeroCnt + 1) * rightZeroCnt / 2;
            }
            cnt += (leftZeroCnt + 1) * (rightZeroCnt + 1);
            prev = onePoses[i];
        }
        return cnt;
    }

    public int numSubarraysWithSum(int[] nums, int goal) {
        int sum = 0;
        Map<Integer, Integer> cnt = new HashMap<>();
        int ret = 0;
        for (int num : nums) {
            cnt.put(sum, cnt.getOrDefault(sum, 0) + 1);
            sum += num;
            ret += cnt.getOrDefault(sum - goal, 0);
        }
        return ret;
    }

    public static void main(String[] args) {
        System.out.println(new T20210708N1().numSubarraysWithSum(new int[] {0,0,0,0,0}, 0));
    }
}
